# -*- coding: utf-8 -*-
# @Author: Cursor
# @Date: 2025-01-08
# @Last Modified by: Tim Liu
# @Last Modified time: 2024-01-08

# @Description: Converts a schema JSON to LLMGraphTransformer compatible tuple

from typing import List, Dict, Tuple

def convert_schema_to_llm_transformer(schema: Dict) -> Tuple[List[str], List[Tuple[str, str, str]], List[str]]:
    """
    Converts a schema JSON to LLMGraphTransformer compatible allowed_nodes, allowed_relationships, and node_properties.
    
    # Sample schema JSON
    schema_json = {
        "node_types": {
            "Author": {
                "properties": {
                    "name": "STRING",
                    "email": "STRING",
                    "bio": "STRING"
                }
            },
            "Blog": {
                "properties": {
                    "title": "STRING",
                    "url": "STRING",
                    "content": "STRING",
                    "published": "DATETIME"
                }
            },
            "Tag": {
                "properties": {
                    "name": "STRING"
                }
            }
        },
        "relationships": {
            "AUTHORED": {
                "source": "Author",
                "target": "Blog",
                "properties": {}
            },
            "HAS_TAG": {
                "source": "Blog",
                "target": "Tag",
                "properties": {}
            }
        }
    }
    
    Args:
        schema (Dict): The schema JSON.

    Returns:
        Tuple[List[str], List[Tuple[str, str, str]], List[str]]: allowed_nodes, allowed_relationships, node_properties.
    """
    allowed_nodes = list(schema["node_types"].keys())
    allowed_relationships = []
    node_properties = set()

    for rel_type, rel_info in schema["relationships"].items():
        source = rel_info["source"]
        target = rel_info["target"]
        allowed_relationships.append((source, rel_type, target))

    for node_type, node_info in schema["node_types"].items():
        for prop in node_info["properties"].keys():
            node_properties.add(prop)

    return allowed_nodes, allowed_relationships, list(node_properties)

def convert_llm_transformer_to_schema(allowed_nodes: List[str], allowed_relationships: List[Tuple[str, str, str]], node_properties: List[str]) -> str:
    """
    Converts LLMGraphTransformer compatible allowed_nodes, allowed_relationships, and node_properties back to a schema JSON string.
    
    Args:
        allowed_nodes (List[str]): List of allowed node types.
        allowed_relationships (List[Tuple[str, str, str]]): List of allowed relationships.
        node_properties (List[str]): List of node properties.

    Returns:
        str: The schema JSON string.
    """
    schema = {
        "node_types": {node: {"properties": {prop: "STRING" for prop in node_properties}} for node in allowed_nodes},
        "relationships": {rel[1]: {"source": rel[0], "target": rel[2], "properties": {}} for rel in allowed_relationships}
    }
    
    return schema